from PySide2.QtWidgets import *
from PySide2.QtCore import *
from PySide2.QtGui import *
import os, time, copy, psutil
import win32gui, win32api, win32process, win32con, win32ui
import numpy as np
from PIL import Image
from Widgets import Label
from Widgets import blueStyle, darkStyle


class GetHwndThread(QThread):  # 获取窗口句柄
    hwndList = Signal(list)  # 获取的句柄列表

    def __init__(self):
        super(GetHwndThread, self).__init__()

    def get_all_hwnd(self, hwnd, mouse):
        if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):  # 判断窗体有效
            windowTitle = win32gui.GetWindowText(hwnd)
            if windowTitle:
                threadID, processID = win32process.GetWindowThreadProcessId(hwnd)  # 获取pid
                processName = psutil.Process(processID).name()  # 获取进程名 比如Photoshop.exe
                self.hwnds[hwnd] = '%s - %s' % (processName, windowTitle)  # 记录进程名和标题名

    def run(self):
        self.hwnds = {}
        win32gui.EnumWindows(self.get_all_hwnd, 0)  # 获取当前有效的窗口
        hwnds = []
        title = []
        hwndsCopy = copy.deepcopy(self.hwnds)  # 深拷贝用于防止pop item后序号错乱
        for hwnd, v in self.hwnds.items():
            for perticular in ['photoshop', 'sai', 'clip']:  # 包含这些关键词的进程优先靠前排列
                if perticular in v.lower():
                    hwnds.append(hwnd)
                    title.append(hwndsCopy.pop(hwnd))
        hwnds += list(hwndsCopy.keys())  # 添加剩余的进程
        title += list(hwndsCopy.values())
        self.hwndList.emit([hwnds, title])  # 发送句柄列表


class AutoSaver(QWidget):  # 自动保存插件
    def __init__(self, saveConfig):
        super(AutoSaver, self).__init__()
        # self.setWindowFlags(Qt.WindowStaysOnTopHint)
        self.setWindowTitle('保存助手')
        self.resize(420, 240)
        layout = QGridLayout(self)
        layout.addWidget(Label('选择程序'), 0, 0, 1, 1)
        self.hwnds = []  # 存放句柄的列表
        self.selectHwnd = 0  # 当前选中的句柄
        self.selectTitle = ''  # 当前选中的进程名
        self.processCombox = QComboBox()  # 进程下拉菜单
        self.processCombox.setMinimumWidth(300)
        self.processCombox.currentIndexChanged.connect(self.selectProcess)
        layout.addWidget(self.processCombox, 0, 1, 1, 3)
        layout.addWidget(Label('保存间隔'), 1, 0, 1, 1)
        self.saveInterval = 15  # 默认15秒保存一次
        self.intervalList = [15 * x for x in range(1, 4)] + [60 * x for x in range(1, 11)]
        self.intervalCombox = QComboBox()
        self.intervalCombox.addItems(['%s 秒' % (x * 15) for x in range(1, 4)] + ['%s 分钟' % x for x in range(1, 11)])
        self.intervalCombox.currentIndexChanged.connect(self.selectInterval)
        layout.addWidget(self.intervalCombox, 1, 1, 1, 1)

        self.saveConfig = saveConfig  # 使用可改变对象的列表来同步主线程里的config
        self.saveCaptureButton = QPushButton('同时截图')
        self.saveCaptureButton.clicked.connect(self.selectCapture)
        layout.addWidget(self.saveCaptureButton, 2, 0, 1, 1)
        self.saveCaptureLabel = QLabel()
        self.saveCaptureLabel.setTextInteractionFlags(Qt.TextSelectableByMouse)
        layout.addWidget(self.saveCaptureLabel, 2, 1, 1, 3)
        if saveConfig[0] and os.path.exists(saveConfig[0]):  # 本地配置里已设置保存路径
            self.captureToken = True
            self.saveCaptureButton.setStyleSheet(blueStyle)
            self.saveCaptureLabel.setText('保存至 ' + saveConfig[0])
        else:
            self.saveCaptureButton.setStyleSheet(darkStyle)
            self.captureToken = False

        self.startToken = False
        self.startButton = QPushButton('开始')
        self.startButton.setFixedHeight(30)
        self.startButton.clicked.connect(self.startAutoSaver)
        layout.addWidget(self.startButton, 3, 0, 1, 4)

        self.refreshCnt = 0
        self.autoSaveTimer = QTimer()
        self.autoSaveTimer.timeout.connect(self.postSaveMessage)

        self.getHwndThread = GetHwndThread()  # 获取当前显示状态窗口句柄的thread
        self.getHwndThread.hwndList.connect(self.refreshProcessCombox)

    def refreshProcessCombox(self, hwndList):
        self.hwnds, self.title = hwndList
        self.selectHwnd = self.hwnds[0]  # 默认选择第一个
        self.selectTitle = self.title[0].replace('.exe - ', '_')
        for s in ['/', '\\', ':', '*', '"', '<', '>', '|', '？']:  # 移除不支持的字符
            self.selectTitle = self.selectTitle.replace(s, '')
        try:
            self.processCombox.currentIndexChanged.disconnect(self.selectProcess)
        except:
            pass
        self.processCombox.clear()
        self.processCombox.addItems(self.title)
        self.processCombox.currentIndexChanged.connect(self.selectProcess)

    def selectProcess(self, index):
        self.selectHwnd = self.hwnds[index]
        self.selectTitle = self.title[index].replace('.exe - ', '_')
        for s in ['/', '\\', ':', '*', '"', '<', '>', '|', '？']:  # 移除不支持的字符
            self.selectTitle = self.selectTitle.replace(s, '')

    def selectInterval(self, index):
        self.saveInterval = self.intervalList[index]

    def selectCapture(self):  # 响应用户点击
        if not self.captureToken:
            path = QFileDialog.getExistingDirectory(self, '选择截图保存路径')
            if path:  # 选择有效路径
                self.saveConfig[0] = path
                self.saveCaptureButton.setStyleSheet(blueStyle)
                self.saveCaptureLabel.setText('保存至 ' + path)  # 显示用户选择路径
                self.captureToken = not self.captureToken
        else:
            self.saveConfig[0] = ''
            self.saveCaptureButton.setStyleSheet(darkStyle)
            self.saveCaptureLabel.clear()  # 清除显示
            self.captureToken = not self.captureToken

    def startAutoSaver(self):
        self.startToken = not self.startToken
        if self.startToken:
            self.startButton.setStyleSheet(blueStyle)
            self.startButton.setText('停止(保存中)')
            self.refreshCnt = self.saveInterval - 1  # 设置初始计时秒数让程序在开始按钮点下后立刻保存一次
            self.postSaveMessage()  # 发送保存信号
            self.processCombox.setEnabled(False)  # 冻结控件
            self.intervalCombox.setEnabled(False)
            self.saveCaptureButton.setEnabled(False)
            self.autoSaveTimer.start(1000)  # 每秒刷新一次
        else:
            self.startButton.setStyleSheet(darkStyle)
            self.startButton.setText('开始')
            self.processCombox.setEnabled(True)  # 控件
            self.intervalCombox.setEnabled(True)
            self.saveCaptureButton.setEnabled(True)
            self.autoSaveTimer.stop()
            self.refreshCnt = 0

    def postSaveMessage(self):
        try:
            self.refreshCnt += 1
            if self.refreshCnt >= self.saveInterval:  # 计时达到设置的间隔秒数
                self.startButton.setText('停止(保存中)')
                self.refreshCnt = 0  # 清零计时
                win32api.keybd_event(17, 0, 0, 0)  # 我也不知道为啥py3的pywin32必须改用keybd_event按住Ctrl键 懵逼
                win32api.keybd_event(83, 0, 0, 0)  # keybd_event按住S键
                win32api.PostMessage(self.selectHwnd, win32con.WM_KEYDOWN, 0x53, 1)  # 按下 S键
                time.sleep(0.01)  # 还得加个延时 不然有的软件反应不过来

                win32api.keybd_event(83, 0, win32con.KEYEVENTF_KEYUP, 0)  # 释放S键
                win32api.keybd_event(17, 0, win32con.KEYEVENTF_KEYUP, 0)  # 释放Ctrl键
                win32api.PostMessage(self.selectHwnd, win32con.WM_KEYUP, 0x53, 1)  # 释放 S键

                if self.captureToken:  # 如果设置了自动保存
                    left, top, right, bot = win32gui.GetWindowRect(self.selectHwnd)
                    w = right - left
                    h = bot - top
                    wDC = win32gui.GetWindowDC(self.selectHwnd)
                    dcObj = win32ui.CreateDCFromHandle(wDC)
                    cDC = dcObj.CreateCompatibleDC()
                    dataBitMap = win32ui.CreateBitmap()
                    dataBitMap.CreateCompatibleBitmap(dcObj, w, h)
                    cDC.SelectObject(dataBitMap)
                    cDC.BitBlt((0, 0), (w, h), dcObj, (0, 0), win32con.SRCCOPY)
                    img = np.frombuffer(dataBitMap.GetBitmapBits(True), dtype='uint8')
                    img.shape = (h, w, 4)
                    img = Image.fromarray(img.astype('uint8')).convert('RGB')
                    b, g, r = img.split()  # 调换颜色通道
                    img = Image.merge('RGB', (r, g, b))
                    timeStamp = time.strftime('%H-%M-%S', time.localtime(time.time()))
                    jpgPath = '%s/%s_%s.jpg' % (self.saveConfig[0], self.selectTitle, timeStamp)
                    img.save(jpgPath, quality=90)
                    dcObj.DeleteDC()
                    cDC.DeleteDC()
                    win32gui.ReleaseDC(self.selectHwnd, wDC)
                    win32gui.DeleteObject(dataBitMap.GetHandle())
            else:
                self.startButton.setText('停止(%s秒后自动保存)' % (self.saveInterval - self.refreshCnt))
        except Exception as e:
            print(str(e))
